from torch import nn
from typing import Tuple


class NiN(nn.Module):
    def __init__(self):
        super(NiN, self).__init__()
        self._max_pool = nn.MaxPool2d(3, stride=2)
        self._flatten = nn.Flatten()
        self._avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self._block1 = self.get_block(1, 96, (11,), (4,))
        self._block2 = self.get_block(96, 256, (5,), (1,), padding=2)
        self._block3 = self.get_block(256, 384, (3,), (1,), padding=1)
        self._block4 = self.get_block(384, 10, (3,), (1,), padding=1)

    def forward(self, x):
        x = self._block1(x)
        x = self._max_pool(x)
        x = self._block2(x)
        x = self._max_pool(x)
        x = self._block3(x)
        x = self._max_pool(x)
        x = self._block4(x)
        x = self._avg_pool(x)
        return self._flatten(x)

    @staticmethod
    def get_block(in_channel: int, out_channel: int, kernel_size: Tuple[int, ...], strides: Tuple[int, ...] = (1,), padding: int = 0):
        return nn.Sequential(
            nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=kernel_size, stride=strides, padding=padding),
            nn.ReLU(),
            nn.Conv2d(in_channels=out_channel, out_channels=out_channel, kernel_size=(1, 1)), nn.ReLU(),
            nn.Conv2d(in_channels=out_channel, out_channels=out_channel, kernel_size=(1, 1)), nn.ReLU()
        )
